//	Aviv Charikar	Yael Tzirulnikov	305343600	307955609

package modelTraining;

import java.io.*;
import java.util.*;

public class BasicModelTraining {
	private HashMap<String, Integer> wordCounts;
	private int countAll;
	
	public BasicModelTraining(String devFilename) throws IOException {
		wordCounts = new HashMap<String, Integer>();
		countAll = 0;
		readDevFile(devFilename);
	}
	
	private void readDevFile(String devFilename) throws IOException {
		BufferedReader reader = new BufferedReader(new FileReader(devFilename));
		String line;
		while ((line = reader.readLine()) != null) {
			// Handle only sentences
			if ((line.length() == 0) || (line.startsWith("<TRAIN")))
				continue;
			
			// Split the words in the sentence and add them to the counts
			String[] lineWords = line.split(" ");
			for (String word : lineWords) {
				if (word.length() > 0) {
					Integer count = wordCounts.get(word);
					if (count == null)
						wordCounts.put(word, 1);
					else
						wordCounts.put(word, count + 1);
					countAll++;
				}
			}
		}
		reader.close();
	}
	
	public Set<String> getAllWords() {
		return wordCounts.keySet();
	}
	
	public int getWordCount(String word) {
		// Return the count of the word or 0 if the word is unseen
		Integer count = wordCounts.get(word);
		if (count == null)
			return 0;
		return count;
	}
	
	public int getCountAll() {
		return countAll;
	}
}
